library(DeclareDesign)
library(tidyverse)
library(readxl)
library(haven)

setwd("~/Dropbox/Fear and coordination/11_JEPS_Replication_Files/")


## set initial parameters for simulations

N <- 100
beta = 0.5
beta_action = 0.1

## calculate noise on self reported fear and alpha amylase in control group

amylase <- read_excel("01_Data/New data/survey with alpha amylase.xlsx")
dat <- read_dta("01_Data/New data/dat3.dta")


dat <- left_join(dat, amylase, by = 'subject_id')

dat$amylasediff <- dat$amylasetreatment - dat$amylasebaseline

fear_mean <- mean(dat$fear.x[dat$Treatment==0], na.rm=T)
fear_sd <- sd(dat$fear.x[dat$Treatment==0], na.rm=T)
amylasediff_mean <- mean(dat$amylasediff[dat$Treatment==0], na.rm=T)
amylasediff_sd <- sd(dat$amylasediff[dat$Treatment==0], na.rm=T)
action_mean <- mean(dat$action_i[dat$Treatment==0], na.rm=T)
action_sd <- sd(dat$action_i[dat$Treatment==0], na.rm=T)


## simulate data

design <- 
  
  declare_model(
    N = N,                # sample size -- we will allow this to grow
    e_fear = rnorm(N, mean = fear_mean, sd = fear_sd), # fear error
    e_amylase = rnorm(N, mean = amylasediff_mean, sd = amylasediff_sd), # amylase error
    e_action = rnorm(N, mean = action_mean, sd = action_sd), # action error
    Z = rbinom(N, size = 1, prob = 0.5), # treatment assignment 
    potential_outcomes(Y_amylase ~ beta*amylasediff_mean*Z + e_amylase),
    potential_outcomes(Y_fear ~ beta*fear_mean*Z + e_fear),
    potential_outcomes(Y_action ~ beta_action*Z + e_action)
  ) +
  
  # reveal outcomes on the basis of realized treatment
  declare_measurement(Y_amylase = reveal_outcomes(Y_amylase ~ Z)) +
  declare_measurement(Y_fear = reveal_outcomes(Y_fear ~ Z)) +
  declare_measurement(Y_action = reveal_outcomes(Y_action ~ Z)) +
  
  # true effect is 0
  declare_inquiry(ATE = beta) +
  
  # use difference-in-means to estimate effect
  declare_estimator(Y_amylase ~ Z, model = difference_in_means, inquiry = "ATE", label = "ATE_amylase") +
  declare_estimator(Y_fear ~ Z, model = difference_in_means, inquiry = "ATE", label = "ATE_fear") +
declare_estimator(Y_action ~ Z, model = difference_in_means, inquiry = "ATE", label = "ATE_action") 

dat <- draw_data(design)


# we want to know about bias and power

diags <- declare_diagnosands(bias = mean(estimate - estimand),
                             power = mean(p.value <= 0.05))


## redesign sample size btw 100 and 1000

designs <- redesign(design, N = seq(100, 1000, by = 100))

diagnosis <- diagnose_design(designs, diagnosands = diags)


## plot fear diagnosands

get_simulations(diagnosis) %>%
  filter(estimator == 'ATE_fear') %>%
  ggplot(aes(x = N, y = estimate)) +
  geom_point(alpha = 0.2) +
  geom_hline(yintercept = beta*fear_mean, color = "red") +
  geom_text(data = data.frame(N = 800, estimate = 0.35, 
                              label = paste0("True ATE = ", round(beta*fear_mean, 2))),
            aes(label = label)) +
  coord_cartesian(ylim = c(-0.1, 1.5)) +
  theme_bw() +
  ggtitle("Estimated ATEs on fear")

ggsave("03_Plots/Figure8.pdf")
get_diagnosands(diagnosis) %>%
  filter(estimator == 'ATE_fear') %>%
  ggplot(aes(x = N, y = power)) +
  geom_point(alpha = 0.5) + geom_line() +
  geom_hline(yintercept = 0.8, color = 'red', lty = 2) +
  theme_bw() +
  ylim(c(0, 1)) +
  theme(text = element_text(size = 20)) +   
  ggtitle("Power on fear")
dev.off()

## plot amylase diagnosands

get_simulations(diagnosis) %>%
  filter(estimator == 'ATE_amylase') %>%
  ggplot(aes(x = N, y = estimate)) +
  geom_point(alpha = 0.2) +
  geom_hline(yintercept = beta*amylasediff_mean, color = "red") +
  geom_text(data = data.frame(N = 800, estimate = 0.5, 
                              label = paste0("True ATE = ", round(beta*amylasediff_mean, 2))),
            aes(label = label)) +
  coord_cartesian(ylim = c(-30, 30)) +
  theme_bw() +
  ggtitle("Estimated ATEs on alpha-amylase")

get_diagnosands(diagnosis) %>%
  filter(estimator == 'ATE_amylase') %>%
  ggplot(aes(x = N, y = power)) +
  geom_point(alpha = 0.5) + geom_line() +
  geom_hline(yintercept = 0.8, color = 'red', lty = 2) +
  theme_bw() +
  ylim(c(0, 1)) +
  theme(text = element_text(size = 20)) +   
  ggtitle("Power on alpha-amylase")


## plot action diagnosands

designs <- redesign(design, beta_action = seq(0.02, 0.4, by = 0.02), N = 432)

diagnosis <- diagnose_design(designs, diagnosands = diags)

get_simulations(diagnosis) %>%
  filter(estimator == 'ATE_action') %>%
  ggplot(aes(x = N, y = estimate)) +
  geom_point(alpha = 0.2) +
  geom_hline(yintercept = beta*action_mean, color = "red") +
  geom_text(data = data.frame(N = 800, estimate = 0.5, 
                              label = paste0("True ATE = ", round(beta*action_mean, 2))),
            aes(label = label)) +
  coord_cartesian(ylim = c(-30, 30)) +
  theme_bw() +
  ggtitle("Estimated ATEs on action")

ggsave("03_Plots/Figure9.pdf")
get_diagnosands(diagnosis) %>%
  filter(estimator == 'ATE_action') %>%
  ggplot(aes(x = beta_action, y = power)) +
  geom_point(alpha = 0.5) + geom_line() +
  geom_hline(yintercept = 0.8, color = 'red', lty = 2) +
  theme_bw() +
  ylim(c(0, 1)) +
  theme(text = element_text(size = 20)) +   
  ggtitle("Power on action")
dev.off()

